from torch.utils.data import Dataset
import torch


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class TextClassificationDataset(Dataset):
    def __init__(self, texts, labels, tokenizer):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]
        inputs = self.tokenizer(text, padding="max_length", max_length=128, truncation=True, return_tensors='pt')
        return {
            'input_ids': inputs['input_ids'].flatten(),
            'attention_mask': inputs['attention_mask'].flatten(),
            'labels': torch.tensor(label, dtype=torch.long),
        }


def get_dataset(texts, labels, tokenizer):
    dataset = TextClassificationDataset(texts, labels, tokenizer)
    return dataset




